import chromadb
from chromadb.utils import embedding_functions

# 初始化客户端，指定持久化目录（不存在会自动创建）
client = chromadb.PersistentClient(path="./chroma_data")  # 数据将保存在当前目录的 chroma_data 文件夹

# 可选：使用默认的嵌入模型（若不指定，需手动提供向量）
default_ef = embedding_functions.DefaultEmbeddingFunction()

# 创建集合（若已存在则获取）
collection = client.get_or_create_collection(
    name="my_collection",  # 集合名称（唯一）
    embedding_function=default_ef  # 绑定嵌入函数（自动将文本转为向量）
)

# # 准备数据
# documents = [
#     "Chroma 是一款轻量级向量数据库",
#     "向量数据库用于存储和检索高维向量",
#     "持久化存储可以避免程序退出后数据丢失"
# ]
# metadatas = [  # 元数据（可选，用于过滤查询）
#     {"source": "doc1"},
#     {"source": "doc2"},
#     {"source": "doc3"}
# ]
# ids = ["id1", "id2", "id3"]  # 每条数据的唯一ID
#
# # 添加数据（自动生成向量并持久化到 ./chroma_data）
# collection.add(
#     documents=documents,
#     metadatas=metadatas,
#     ids=ids
# )

# 搜索与"向量数据库"最相似的文档（返回前2条）
results = collection.query(
    query_texts=["向量数据库"],  # 待查询的文本（自动转为向量）
    n_results=2  # 返回匹配度最高的2条结果
)

print("查询结果：")
for doc, score in zip(results["documents"][0], results["distances"][0]):
    print(f"文档：{doc}，相似度：{score:.4f}")